package com.wave.wavesystem.ai.chatmemory;

import com.wave.common.util.Func;
import com.wave.wavesystem.ai.pojo.entity.AiConversation;
import com.wave.wavesystem.ai.pojo.entity.AiConversationMessage;
import com.wave.wavesystem.ai.service.IAiConversationMessageService;
import com.wave.wavesystem.ai.service.IAiConversationService;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Transactional;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

@Slf4j
@Component
public class WaveChatMemory implements ChatMemory {

    @Resource
    private IAiConversationService conversationService;
    @Resource
    private IAiConversationMessageService messageService;

    // 在 WaveChatMemory 的 add 方法中添加日志
    @Override
    @Transactional(rollbackFor = Exception.class)
    public void add(String conversationId, List<Message> messages) {
        if ("default".equals(conversationId)) {
            return;
        }
        AiConversation conversation = conversationService.queryByConversationId(conversationId);
        // 如果  conversion 为空 创建对话
        if (Func.isEmpty(conversation)) {
            conversation = new AiConversation();
            conversation.setConversationId(conversationId);
            conversation.setTotalMessages(messages.size());
            if (!messages.isEmpty()) {
                String title = messages.getFirst().getText().trim();
                if (title.length() > 50) {
                    conversation.setTitle(title.substring(0, 50));
                } else {
                    conversation.setTitle(title);
                }
            }
            conversationService.save(conversation);
        }
        // 不为空更新对话信息
        else {
            if (Func.isEmpty(conversation.getTitle())) {
                if (!messages.isEmpty()) {
                    String title = messages.getFirst().getText().trim();
                    if (title.length() > 50) {
                        conversation.setTitle(title.substring(0, 50));
                    } else {
                        conversation.setTitle(title);
                    }
                }
                conversationService.updateById(conversation);
            }
        }
        conversationService.updateChain().eq(AiConversation::getId, conversation.getId())
                .setRaw(AiConversation::getTotalMessages, "totalMessages + " + messages.size())
                .update();
        log.info("开始保存对话记忆，会话ID: {}, 消息数量: {}", conversationId, messages.size());
        if (Func.isEmpty(messages)) {
            log.info("消息列表为空，跳过保存");
            return;
        }

        // 每条消息 message 开始保存到数据库中
        try {
            List<AiConversationMessage> myMessages = messages.stream().map(message -> {
                log.info("转换消息: role={}, content={}", message.getMessageType().getValue(), message.getText());
                return AiConversationMessage.builder().conversationId(conversationId)
                        .role(message.getMessageType().getValue()).content(message.getText()).build();
            }).toList();

            boolean saveResult = messageService.saveBatch(myMessages);
            log.info("保存对话记忆结果: {}, 会话ID: {}", saveResult, conversationId);

        } catch (Exception e) {
            log.info("保存对话记忆异常", e);
            throw e;
        }
    }

    @Override
    public List<Message> get(String conversationId) {
        if ("default".equals(conversationId)){
            return Collections.emptyList();
        }
        if (Func.isEmpty(conversationId)) {
            return Collections.emptyList();
        }
        List<Message> returnMessages = new ArrayList<>();
        List<AiConversationMessage> messages = messageService.listByConversationId(conversationId);
        if (Func.isEmpty(messages)) {
            return List.of();
        }
        for (AiConversationMessage message : messages) {
            MessageType messageType = MessageType.fromValue(message.getRole());
            switch (messageType) {
                case MessageType.ASSISTANT:
                    returnMessages.add(new AssistantMessage(message.getContent()));
                    break;
                case MessageType.USER:
                    returnMessages.add(new UserMessage(message.getContent()));
                    break;
                default:
                    return Collections.emptyList();
            }
        }
        return returnMessages;
    }

    @Override
    public void clear(String conversationId) {
        conversationService.clear(conversationId);
    }

}
